"""
Unit tests for modules in the lib.ml package.

Copyright (C) 2013 Jerrad Michael Genson

This program is free software: you can redistribute it and/or modify
it under the terms of the BSD 3-Clause License as published by
the Open Source Initiative.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
BSD 3-Clause License for more details.

You should have received a copy of the BSD 3-Clause License
along with this program.  If not, see
<https://betelgeuse-ml.googlecode.com/hg/LICENSE>

"""

# Python build-in modules
import unittest
import random
import cProfile

# Third-party modules
import numpy as np

# Betelgeuse modules
import optimize
import ga
import pso
import sa

# Module header
__author__ = 'Jerrad Genson'
__contact__ = 'jerradgenson@gmail.com'
__copyright__ = 'Copyright 2013 Jerrad Michael Genson'
__license__ = 'BSD 3-Clause'


def bubble_pass(numbers):
    """ Run one pass of bubble sort on a number array. """

    for index in range(len(numbers) - 1):
        if numbers[index] > numbers[index + 1]:
            numbers = optimize.swap(index, index + 1, numbers)

    return numbers


def bubble_optimize(number_lists):
    """ Run one pass of bubble sort on a list of unsorted number lists. """

    return [bubble_pass(numbers) for numbers in number_lists]


def list_quality(numbers):
    """ Measure the sortedness of a number list or array. """

    unsorted = 0
    for index in range(len(numbers) - 1):
        if numbers[index] <= numbers[index + 1]:
            unsorted += 1.0

    return unsorted / (len(numbers) - 1)


class TestOptimizeFunctions(unittest.TestCase):
    """ Unit tests for optimize.py functions """

    def setUp(self):
        self.array1 = np.array([1, 2, 3, 4])
        self.array2 = np.array([1, 2, 3, 4])
        self.unsorted = np.array([4, 1, 3, 2])
        self.array_list1 = [self.array2, self.array1, self.array2]
        self.array_list2 = [self.array2, self.array2, self.array2]
        self.array_list3 = []

    def test_create_solutions(self):
        """ test create_solutions """

        array = np.array([5, 8, 3])
        random.seed(1101)
        states = optimize.create_solutions(array, 2, ga.create_shuffle_mutate())
        self.assertEqual(len(states), 2)
        for index in range(2):
            self.assertTrue(5 in states[index])
            self.assertTrue(8 in states[index])
            self.assertTrue(3 in states[index])

    def test_combine_optimizers(self):
        """ test combine_optimizers """

        ga_optimize = ga.create_optimizer(list_quality,
                                          ga.create_shuffle_mutate(),
                                          ga.create_point_crossover(),
                                          select=ga.create_truncate_select())

        composite_optimize = optimize.combine_optimizers((ga_optimize,
                                                          bubble_optimize),
                                                         (1, 2))

        solution = composite_optimize([self.unsorted])
        self.assertTrue((solution == np.array([1, 2, 3, 4])).all())

    def test_conditional_optimize(self):
        """ test conditional_optimize """

        solution1 = optimize.conditional_optimize((x for x in [self.unsorted]),
                                                  bubble_optimize,
                                                  max_iterations=1)

        solution2 = optimize.conditional_optimize((x for x in [self.unsorted]),
                                                  bubble_optimize,
                                                  min_quality=1,
                                                  check_quality=list_quality)

        self.assertTrue((solution1 == np.array([1, 3, 2, 4])).all())
        self.assertTrue((solution2 == np.array([1, 2, 3, 4])).all())

    def test_find_best(self):
        """ test find_best function """

        states = [[2, 3, 1],
                  [1, 2, 3],
                  [3, 1, 2]]

        best = optimize.find_best(states, list_quality)
        self.assertEqual(best[0], [1, 2, 3])
        self.assertEqual(best[1], 1.0)

    def test_swap(self):
        """ test swap function """

        new_array = optimize.swap(1, 3, self.array1)
        self.assertTrue((self.array1 == np.array([1, 2, 3, 4])).all())
        self.assertTrue((new_array == np.array([1, 4, 3, 2])).all())

    def test_array_in_arrays(self):
        """ test array_in_list function """

        self.assertTrue(optimize.array_in_arrays(self.array1, self.array_list1))
        self.assertFalse(optimize.array_in_arrays(self.array1, self.array_list2))
        self.assertFalse(optimize.array_in_arrays(self.array1, self.array_list3))

    def test_remove_array(self):
        """ test remove_array """

        new_list1 = optimize.remove_array(self.array1, self.array_list1)
        new_list2 = optimize.remove_array(self.array1, self.array_list2)
        new_list3 = optimize.remove_array(self.array1, self.array_list3)
        self.assertEqual(len(new_list1), 2)
        self.assertEqual(len(new_list2), 3)
        self.assertEqual(len(new_list3), 0)
        self.assertTrue(new_list1[0] is self.array2)
        self.assertTrue(new_list1[1] is self.array2)


class TestGAFunctions(unittest.TestCase):
    """ unit tests for ga.py """

    def setUp(self):
        self.intparent = np.array([1, 3, 5, 7, 9, 0, 2, 4, 6, 8])
        self.intparent2 = np.array([0, 2, 4, 6, 8, 1, 3, 5, 7, 9])
        self.intparent3 = np.array([11, 12, 13, 14])
        self.realparent = np.array([1.0, 2.5, 3.5, 4.5, 5.0])
        self.sorted = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
        self.chromosomes = [self.intparent, self.intparent2, self.intparent3]

    def test_generic_mutation(self):
        """ test create_generic_mutation function """

        # Test generic_mutate function without bounds checking.
        mutate = ga.create_generic_mutate()
        random.seed(10)
        child1 = mutate(self.intparent, 0.2)
        child2 = mutate(self.intparent, 0.25)
        child3 = mutate(self.intparent, 0.3)
        self.assertTrue((self.intparent == np.array([1, 3, 5, 7, 9, 0, 2, 4, 6, 8])).all())
        self.assertTrue((child1 != self.intparent).any())
        self.assertTrue((child2 != self.intparent).any())
        self.assertTrue((child3 != self.intparent).any())

        # Test generic_mutate function with bounds checking.
        mutate = ga.create_generic_mutate((0, 5))
        random.seed(10)
        child1 = mutate(self.realparent, 0.5)
        child2 = mutate(self.realparent, 0.75)
        child3 = mutate(self.realparent, 1)
        self.assertTrue((self.realparent == np.array([1.0, 2.5, 3.5, 4.5, 5.0])).all())
        self.assertTrue((child1 != self.realparent).any())
        self.assertTrue((child1 >= 0).all())
        self.assertTrue((child1 <= 5).all())
        self.assertTrue((child2 != self.realparent).any())
        self.assertTrue((child2 >= 0).all())
        self.assertTrue((child2 <= 5).all())
        self.assertTrue((child3 != self.realparent).any())
        self.assertTrue((child3 >= 0).all())
        self.assertTrue((child3 <= 5).all())

    def test_shuffle_mutate(self):
        """ test shuffle_mutate function """

        random.seed(1010)
        child1 = ga.create_shuffle_mutate()(self.intparent, 0.5)
        child2 = ga.create_shuffle_mutate()(self.intparent, 0.75)
        child3 = ga.create_shuffle_mutate()(self.intparent, 1.0)
        self.assertTrue((self.intparent == np.array([1, 3, 5, 7, 9, 0, 2, 4, 6, 8])).all())
        self.assertTrue((child1 != self.intparent).any())
        self.assertTrue((child2 != self.intparent).any())
        self.assertTrue((child3 != self.intparent).any())
        for integer in self.intparent:
            self.assertTrue(integer in child1)
            self.assertTrue(integer in child2)
            self.assertTrue(integer in child3)

    def test_point_crossover(self):
        """ test point_crossover function """

        child1 = ga.create_point_crossover()(self.intparent, self.intparent2)
        child2 = ga.create_point_crossover()(self.intparent2, self.intparent)
        self.assertTrue((child1 == np.array([1, 3, 5, 7, 9, 1, 3, 5, 7, 9])).all())
        self.assertTrue((child2 == np.array([0, 2, 4, 6, 8, 0, 2, 4, 6, 8])).all())

    def test_random_crossover(self):
        """ test random_crossover function """

        random.seed(10)
        child1 = ga.create_random_crossover()(self.intparent, self.intparent2)
        child2 = ga.create_random_crossover()(self.intparent, self.intparent2)
        child3 = ga.create_random_crossover()(self.intparent2, self.intparent)
        self.assertTrue((child1 != self.intparent).any())
        self.assertTrue((child1 != self.intparent2).any())
        self.assertTrue((child2 != self.intparent).any())
        self.assertTrue((child2 != self.intparent2).any())
        self.assertTrue((child3 != self.intparent).any())
        self.assertTrue((child3 != self.intparent2).any())
        for index in range(len(self.intparent)):
            self.assertTrue(child1[index] == self.intparent[index] or
                            child1[index] == self.intparent2[index])

            self.assertTrue(child2[index] == self.intparent[index] or
                            child2[index] == self.intparent2[index])

            self.assertTrue(child3[index] == self.intparent[index] or
                            child3[index] == self.intparent2[index])

    def test_tournament_select(self):
        """ test tournament_select function """

        survivor = ga.create_tournament_select()(self.chromosomes, list_quality, 1)
        self.assertEqual(len(survivor), 1)
        self.assertTrue((survivor[0] == self.intparent3).all())

    def test_roulette_select(self):
        """ test roulette_select function """

        random.seed(10)
        survivor = ga.create_roulette_select()(self.chromosomes, list_quality, 1)
        self.assertEqual(len(survivor), 1)
        self.assertTrue((survivor[0] == self.chromosomes[0]).all() or
                        (survivor[0] == self.chromosomes[1]).all()
                            (survivor[0] == self.chromosomes[2]).all())

    def test_truncate_select(self):
        """ test truncate_select function """

        survivor = ga.create_truncate_select()(self.chromosomes, list_quality, 1)
        self.assertEqual(len(survivor), 1)
        self.assertTrue((survivor[0] == self.intparent3).all())

    def test_create_adaptor(self):
        """ test create_adaptor function """

        adapt = ga.create_adaptor(list_quality)
        mut_size, pop_size = adapt(self.chromosomes, 1)
        self.assertAlmostEqual(mut_size, 0.1, places=1)
        self.assertEqual(pop_size, ga.MIN_POPULATION_SIZE)

    def test_create_optimizer(self):
        """ test genetic algorithm optimizer """

        ga_optimize = ga.create_optimizer(list_quality,
                                          ga.create_shuffle_mutate(),
                                          ga.create_point_crossover(),
                                          select=ga.create_truncate_select(),
                                          elitism=True,
                                          mutation_rate=1)

        # Fully optimize problem using GA.
        # TODO: Do we need to seed random here?
        random.seed(10)
        chromosomes = optimize.create_solutions(np.array([4, 1, 3, 2]),
                                                25,
                                                ga.create_shuffle_mutate())

        solutions = optimize.conditional_optimize(chromosomes,
                                                  ga_optimize,
                                                  min_quality=1,
                                                  check_quality=list_quality)

        best = optimize.find_best(solutions, list_quality)
        self.assertTrue((best[0] == np.array([1, 2, 3, 4])).all())

    def test_group_chromosomes(self):
        """ test group_chromosomes function """

        self.chromosomes.extend([self.intparent, self.intparent])
        self.chromosomes.append(self.intparent2)
        groups = ga.group_chromosomes(self.chromosomes, list_quality)
        self.assertEqual(len(groups), 3)
        # Tests for first ChromosomeGroup.
        self.assertTrue(optimize.array_value_equal(groups[1].chromosome,
                                                   self.intparent))

        self.assertAlmostEqual(groups[1].fitness, 0.89, places=1)
        self.assertEqual(groups[1].count, 3)

        # Tests for second ChromosomeGroup.
        self.assertTrue(optimize.array_value_equal(groups[0].chromosome,
                                                   self.intparent2))

        self.assertAlmostEqual(groups[0].fitness, 0.89, places=1)
        self.assertEqual(groups[0].count, 2)

        # Tests for third ChromosomeGroup.
        self.assertTrue(optimize.array_value_equal(groups[2].chromosome,
                                                   self.intparent3))

        self.assertAlmostEqual(groups[2].fitness, 1.0, places=1)
        self.assertEqual(groups[2].count, 1)


class TestPSO(unittest.TestCase):
    """ unit tests for particle swarm optimization """

    def test_update_position(self):
        """ test create_update_position and update_position functions. """

        particle = pso.Particle([4], dtype='int')
        update_position = pso.create_update_position()
        new_position = update_position(particle)
        self.assertTrue(optimize.array_value_equal(particle, new_position))

    # The randomization algorithm changed in Py3k, so all tests that rely on
    # the behavior of the old algorithm will fail.
    def test_update_velocity(self):
        """ test create_update_velocity and update_velocity functions. """

        particle = pso.Particle([3], dtype='int')
        particle[0] = 1
        particle[1] = 2
        particle[2] = 3
        gbest = pso.Particle([3], dtype='int')
        gbest[0] = 4
        gbest[1] = 5
        gbest[2] = 6
        random.seed(10)
        update_velocity = pso.create_update_velocity()
        new_velocity = update_velocity(particle, gbest)
        self.assertAlmostEqual(new_velocity[0], 1.43052914, places=4)
        self.assertAlmostEqual(new_velocity[1], 0.28772395, places=4)
        self.assertAlmostEqual(new_velocity[2], -0.85508124, places=4)

    def test_particle_swarm_optimizer(self):
        """ test create_optimizer and particle_swarm_optimizer """

        update_velocity = pso.create_update_velocity()
        update_position = pso.create_update_position()
        pso_optimize = pso.create_optimizer(list_quality,
                                            update_velocity,
                                            update_position)

        particle = pso.Particle([4], dtype='int')
        particle[0] = 4
        particle[1] = 3
        particle[2] = 2
        particle[3] = 1
        swarm = optimize.create_solutions(particle,
                                          20,
                                          ga.create_shuffle_mutate())

        optimized_swarm = optimize.conditional_optimize(swarm,
                                                        pso_optimize,
                                                        min_quality=1,
                                                        check_quality=list_quality)

        best = optimize.find_best(optimized_swarm, list_quality)
        self.assertEqual(list_quality(best[0]), 1.0)


class TestSimulatedAnnealing(unittest.TestCase):
    """ unit tests for simulated annealing """

    def test_create_optimizer(self):
        """ test create_optimizer and sa_optimize """

        system_state = np.array([4, 3, 2, 1])
        sa_optimize = sa.create_optimizer(lambda x: 1 - list_quality(x),
                                          ga.create_shuffle_mutate(),
                                          100,
                                          0.9,
                                          0.5)

        solution = optimize.conditional_optimize([system_state], sa_optimize,
                                                 min_quality=1,
                                                 check_quality=list_quality)

        self.assertTrue(optimize.array_value_equal(solution, np.array([1, 2, 3, 4])))


if __name__ == '__main__':
    unittest.main()
